namespace MoreLinq.Test.Async;

using System;
using System.Collections.Generic;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;

static partial class TestExtensions
{
    public static IAsyncEnumerable<T> Yield<T>(this IAsyncEnumerable<T> source) =>
        Yield(source, _ => true);

    public static IAsyncEnumerable<T>
        Yield<T>(this IAsyncEnumerable<T> source,
            Func<T, bool> predicate) =>
        Yield(source, shouldEndSynchronously: false, predicate);

    public static IAsyncEnumerable<T>
        Yield<T>(this IAsyncEnumerable<T> source,
            bool shouldEndSynchronously,
            Func<T, bool> predicate)
    {
        if (source is null) throw new ArgumentNullException(nameof(source));
        if (predicate is null) throw new ArgumentNullException(nameof(predicate));

        return Async();

        async IAsyncEnumerable<T> Async([EnumeratorCancellation]CancellationToken cancellationToken = default)
        {
            await foreach (var item in source.WithCancellation(cancellationToken))
            {
                if (predicate(item))
                    await Task.Yield();
                yield return item;
            }

            if (!shouldEndSynchronously)
                await Task.Yield();
        }
    }
}
